package leetcode.problems.P1277统计全为1的正方形子矩阵;

public class CountSquares {
    public static void main(String[] args) {
        int[][] matrix = new int[][]{
                {0, 1, 1, 1},
                {1, 1, 1, 1},
                {0, 1, 1, 1}
        };
        Solution solution = new Solution();
        System.out.println(solution.countSquares(matrix));
    }
}
class Solution {
    int m, n;

    public int countSquares(int[][] matrix) {
        m = matrix.length;
        if (m < 1) {
            return 0;
        }
        n = matrix[0].length;
        MatrixSum sum = new MatrixSum(matrix);
        int ans = sum.query(0, 0, m - 1, n - 1);
        for (int i = 0; i < m; i++) {
            for (int j = 0; j < n; j++) {
                if (matrix[i][j] == 1) {
                    for (int k = 1; k < Math.min(m, n); k++) {
                        if (!valid(i + k, j + k)) {
                            continue;
                        }
                        int squareSum = sum.query(i, j, i + k, j + k);
                        ans += squareSum == (k + 1) * (k + 1) ? 1 : 0;
                    }
                }
            }
        }
        return ans;
    }

    private boolean valid(int i, int j) {
        return i >= 0 && i < m && j >= 0 && j < n;
    }

    private static class MatrixSum {
        int[][] sum;
        int m, n;

        MatrixSum(int[][] matrix) {
            m = matrix.length;
            if (m < 1) {
                return;
            }
            n = matrix[0].length;
            sum = new int[m][n];
            for (int i = 0; i < m; i++) {
                for (int j = 0; j < n; j++) {
                    int area1 = 0, area2 = 0, area3 = 0, area4 = matrix[i][j];
                    if (j > 0) {
                        area1 = sum[i][j - 1];
                    }
                    if (i > 0) {
                        area2 = sum[i - 1][j];
                    }
                    if (i > 0 && j > 0) {
                        area3 = sum[i - 1][j - 1];
                    }
                    sum[i][j] = area1 + area2 - area3 + area4;
                }
            }
        }

        public int query(int r1, int c1, int r2, int c2) {
            int area1 = sum[r2][c2], area2 = 0, area3 = 0, area4 = 0;
            if (r1 > 0) {
                area2 = sum[r1 - 1][c2];
            }
            if (c1 > 0) {
                area3 = sum[r2][c1 - 1];
            }
            if (r1 > 0 && c1 > 0) {
                area4 = sum[r1 - 1][c1 - 1];
            }
            return area1 - area2 - area3 + area4;
        }
    }
}